from Sanae3D.lib.ModelObject import Model3D
from Sanae3D.lib.StructReader import StructReader

#Some obj files group vertices separately from faces, while others group
#them together

class Obj_Object(Model3D):
    '''Obj file format'''
    
    def __init__(self, vgroup=False, obj3D=None, inFile=None, outFile=None):
        
        super(Obj_Object, self).__init__("Obj")
        self.obj3D = obj3D
        self.inFile = inFile
        self.outFile = outFile
        self.vertex_grouping = 0
        
        self.normals = []
        self.texCoords = []
        
    def parse_materials(self, line):
        
        libName = ' '.join(line.split()[1:])
        try:
            matLib = open(libName)
        except:
            print "Couldn't find library: %s" %libName
            return
        matNum = -1  #0-index since the first material will have index 0
        
        line = matLib.readline()
        while line:
            line = line.strip()
            if line.startswith("newmtl"):
                matNum += 1
                matName = line.split()[1:]
                self.create_material(matNum)  
                self.add_material_name(matNum, matName)
            elif line.startswith("Ka"):
                ambient = line.split()[1:]
                self.add_ambient(matNum, ambient)
            elif line.startswith("Kd"):
                diffuse = line.split()[1:]
                self.add_diffuse(matNum, diffuse)
            elif line.startswith("Ks"):
                specular = line.split()[1:]
                self.add_specular(matNum, specular)
            elif line.startswith("d") or line.startswith("Tr"):
                pass
            elif line.startswith("Ns"):
                power = line.split()[1]
                self.add_power(matNum, power)
            elif line.startswith("map"):
                texName = ' '.join(line.split()[1:])
                self.add_texture_name(matNum, texName)
            line = matLib.readline()
        matLib.close()
        
    def set_material(self, line):
        
        matName = ' '.join(line.split()[1:])
        try:
            matNum = self.get_material_number(matName)
            return matNum
        except:
            return None
        
    def parse_vertices(self, objName, line, vertNum):
        
        coords = line.split()[1:]
        
        coords[0] = float(coords[0]) * -1.0
        self.create_vertex(objName, vertNum)
        self.add_coords(objName, vertNum, coords)
        
    def parse_texture_coords(self, objName, line, vertNum):
        
        texCoords = line.split()[1:]
        self.texCoords.append(texCoords)
    
    def parse_normals(self, objName, line, vertNum):
        
        normals = line.split()[1:]
        self.normals.append(normals)
    
    def parse_faces(self, objName, line, faceNum, offsets, matNum=None):
        '''Assumes the normals and texture coords are stored with vertices.
        Use offsets to reference each vertex, texture, or normal relative to
        the object. Remember that obj files use 1-based index'''
        
        faces = line.split()[1:]
        
        verts = []
        texs = []
        norms = []
        
        for face in faces:
            elem = face.split("/")
            vert = int(elem[0]) - offsets[0] - 1
            
            #case v
            verts.append(vert)
            #case v/vt
            if len(elem) == 2:
                uv = int(elem[1]) - offsets[1] - 1
                texs.append(uv)
                norms.append(norm)
                
            #case v/vt/vn
            elif len(elem) == 3:
                
                #case v//vn
                if elem[1]:
                    uv = int(elem[1]) - offsets[1] -1
                    texs.append(uv)
                norm = int(elem[2]) - offsets[2] - 1
                norms.append(norm)
                
        self.create_face(objName, faceNum)
        self.add_face_verts(objName, faceNum, verts)
        
        if texs:
            face_uv = []
            for i in range(len(texs)):
                face_uv.extend(self.texCoords[texs[i]])
            self.add_face_uv(objName, faceNum, face_uv)
            
        if norms:
            face_norms = []
            for i in range(len(norms)):
                face_norms.extend(self.normals[norms[i]])
            self.add_face_normals(objName, faceNum, face_norms)
        
        if matNum != None:
            self.add_face_mat(objName, faceNum, matNum)
            
    def parse_object(self, line):
        
        objName = ' '.join(line.split()[1:])
        if not objName:
            objName = "temp"
        self.create_object(objName)
        return objName
    
    def parse_file(self):
        
        vertNum, texNum, normNum, faceNum = 0, 0, 0, 0
        
        vertOff = 0
        faceOff = 0
        texOff = 0
        normOff = 0
        
        objName = "obj"
        self.create_object(objName)
    
        line = self.inFile.readline()
        while line:
            if line.startswith("g"):
                #uses vertex grouping, can delete temp object
                self.delete_object("obj")
                objName = self.parse_object(line)
                
                #separates vertices from faces
                if objName == 'temp':
                    self.normals = []
                    self.texCoords = []
                    self.vertex_grouping = 1
                    #set offsets
                else:
                    #replace temp name with actual name
                    if self.vertex_grouping:
                        self.rename_object("temp", objName)
                    else:
                        vertOff += vertNum
                        texOff += texNum
                        normOff += normNum
                    #reset count
                    vertNum, texNum, normNum, faceNum = 0, 0, 0, 0
                    
            elif line.startswith("v "):
                self.parse_vertices(objName, line, vertNum)
                vertNum += 1
            elif line.startswith("vt"):
                self.parse_texture_coords(objName, line, texNum)
                texNum += 1
            elif line.startswith("vn"):
                self.parse_normals(objName, line, normNum)
                normNum += 1
            elif line.startswith("f"):
                self.parse_faces(objName, line, faceNum, \
                                 [vertOff, texOff, normOff])
                faceNum += 1
            elif line.startswith("mtllib"):
                self.parse_materials(line)
            elif line.startswith("usemtl"):
                matNum = self.set_material(line)
            line = self.inFile.readline()
        
        self.delete_object("temp")
            
    #Export methods
    
    def print_vertices(self, objName, offset):
        
        vertNum = self.obj3D.vert_count(objName)
        for i in range(0, vertNum):
            coords = self.obj3D.get_coords(objName, i)
            for j in range(0, len(coords)):
                coords[j] = str(coords[j])
            self.outFile.write("v %s\n" %' '.join(coords))
    
    def print_normals(self, objName, faces, offset):
        
        vertNum = self.obj3D.vert_count(objName)
        for i in range(0, vertNum):
            norms = self.obj3D.get_vert_normals(objName, i)
            if norms:
                
                str_norms = ' '.join([str(norms[0]), str(norms[1]), str(norms[2])])
                self.outFile.write("vn %s\n" %''.join(str_norms))
            else:
                self.outFile.write("vn 1.0 1.0 1.0\n")
                
    def print_texture_coords(self, objName, faces, offset):
        
        vertNum = self.obj3D.vert_count(objName)
        for i in range(0, vertNum):
            uv = self.obj3D.get_vert_uv(objName, i)
            if uv:
                str_uv = ' '.join([str(uv[0]), str(uv[1])])
                self.outFile.write("vt %s\n" %str_uv)
            
    def print_faces(self, objName):
        
        faceNum = self.obj3D.face_count(objName)
        print faceNum
        for i in range(0, faceNum):
            verts = self.obj3D.get_face_verts(objName, i)
            v1, v2, v3 = verts[0] + 1, verts[1] + 1, verts[2] + 1
            
            self.outFile.write("f %s//%s " %(v1, v1))
            self.outFile.write("%s//%s " %(v2, v2))
            self.outFile.write("%s//%s\n" %(v3, v3))
    
    def print_materials(self):
        
        pass
    
    def print_objects(self):
        
        faceOff, vertOff, normOff, texOff = 0, 0, 0, 0
        
        objs = self.obj3D.get_objects()
        for objName in objs:
            self.outFile.write("#obj Object\n\n")
            self.outFile.write("g %s\n\n" %objName)
            
            faceList = self.obj3D.get_faces(objName)
            vertOff = self.print_vertices(objName, vertOff)
            texOff = self.print_texture_coords(objName, faceList, texOff)
            normOff = self.print_normals(objName, faceList, normOff)
            self.print_faces(objName)
    
    def print_file(self):
        
        self.outFile.write("#Exported from Sanae\n")
        self.outFile.write("\n")
        self.print_objects()
        self.outFile.close()
        
def read_file(path):
    
    openFile = open(path)
    obj = Obj_Object(inFile=openFile)
    obj.parse_file()
    openFile.close()
    return obj

def write_file(filename, modelObj):
    
    openFile = open(filename + ".obj", 'w')
    obj = Obj_Object(obj3D=modelObj, outFile=openFile)
    obj.print_file()
    openFile.close()
    
def definitions():
    
    return "OBJ", "OBJ", "Wavefront Obj"
    
if __name__ == '__main__':
        
    files = ["boots.obj", "lod1_shield_trollout.obj"]
    obj = read_file(files[0])